package wiki.xsx.core.bytecode.transformer;

import javassist.CtBehavior;
import javassist.CtClass;
import javassist.CtMethod;
import lombok.SneakyThrows;
import wiki.xsx.core.bytecode.BytecodeUtil;
import wiki.xsx.core.bytecode.model.Bytecode;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.function.BiConsumer;

/**
 * 抽象转换器
 *
 * @author xsx
 * @date 2023/1/31
 * @since 1.8
 * <p>
 * Copyright (c) 2023 xsx All Rights Reserved.
 * easy-bytecode is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 * http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 * </p>
 */
public abstract class AbstractTransformer implements ClassFileTransformer {
    /**
     * 字节码缓存
     */
    private static final Map<String, Bytecode> CACHE = new HashMap<>();
    /**
     * 是否写入
     */
    protected Boolean isWrite;
    /**
     * 文件目录
     */
    protected String fileDir;
    /**
     * 全类名
     */
    protected String className;
    /**
     * 方法名
     */
    protected String methodName;

    /**
     * 有参构造
     *
     * @param isWrite    是否写入
     * @param fileDir    文件目录
     * @param className  类名称
     * @param methodName 方法名称
     */
    public AbstractTransformer(Boolean isWrite, String fileDir, String className, String methodName) {
        this.isWrite = Optional.ofNullable(isWrite).orElse(Boolean.FALSE);
        this.fileDir = Optional.ofNullable(fileDir).orElse(".");
        this.className = className;
        this.methodName = methodName;
    }

    /**
     * 执行转换
     *
     * @param ctClass ct类
     * @return 返回字节码
     */
    public abstract byte[] execute(CtClass ctClass);

    /**
     * 转换
     *
     * @param loader              类加载器
     * @param className           类名称
     * @param classBeingRedefined 类型
     * @param protectionDomain    保护域
     * @param classfileBuffer     字节码
     * @return 返回字节码
     */
    @SneakyThrows
    @Override
    public byte[] transform(
            ClassLoader loader,
            String className,
            Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain,
            byte[] classfileBuffer
    ) {
        if (className != null && className.replaceAll("/", ".").intern().equals(this.className)) {
            CtClass ctClass = BytecodeUtil.getClassPool().get(this.className);
            if (ctClass.isFrozen()) {
                ctClass.defrost();
            }
            if (!CACHE.containsKey(this.className)) {
                CACHE.put(this.className, new Bytecode(classBeingRedefined, classfileBuffer));
            }
            byte[] bytes = this.execute(ctClass);
            if (this.isWrite) {
                ctClass.writeFile(this.fileDir);
            }
            return bytes;
        }
        return classfileBuffer;
    }

    /**
     * 获取类型
     *
     * @return 返回类型
     */
    @SneakyThrows
    public Class<?> getClassType() {
        return Class.forName(this.className);
    }

    /**
     * 处理构造方法
     *
     * @param parameterTypes 参数类型
     * @param ctClass        ct类
     * @param function       函数式接口
     * @return 返回字节码
     */
    @SneakyThrows
    protected byte[] processConstructor(List<Class<?>> parameterTypes, CtClass ctClass, ProcessFunction<CtClass, CtBehavior> function) {
        return this.dealWith(ctClass, parameterTypes, (cc, ts) -> this.dealWithConstructor(cc, ts, function));
    }

    /**
     * 处理普通方法
     *
     * @param parameterTypes 参数类型
     * @param ctClass        ct类
     * @param function       函数式接口
     * @return 返回字节码
     */
    @SneakyThrows
    protected byte[] processMethod(List<Class<?>> parameterTypes, CtClass ctClass, ProcessFunction<CtClass, CtBehavior> function) {
        return this.dealWith(ctClass, parameterTypes, (cc, ts) -> this.dealWithMethod(cc, ts, function));
    }

    /**
     * 移除缓存
     *
     * @return 返回字节码对象
     */
    protected Bytecode removeCache() {
        return CACHE.remove(this.className);
    }

    /**
     * 通用方法处理
     *
     * @param ctClass        ct类
     * @param parameterTypes 参数列表
     * @param consumer       函数式接口
     * @return 返回字节码
     */
    @SneakyThrows
    private byte[] dealWith(CtClass ctClass, List<Class<?>> parameterTypes, BiConsumer<CtClass, CtClass[]> consumer) {
        CtClass[] paramTypes = this.getMethodParamTypes(parameterTypes);
        try {
            consumer.accept(ctClass, paramTypes);
        } finally {
            this.detachMethodParams(paramTypes);
        }
        return ctClass.toBytecode();
    }

    /**
     * 处理构造方法
     *
     * @param ctClass    ct类
     * @param paramTypes 参数列表
     * @param function   函数式接口
     */
    @SneakyThrows
    private void dealWithConstructor(CtClass ctClass, CtClass[] paramTypes, ProcessFunction<CtClass, CtBehavior> function) {
        function.process(ctClass, ctClass.getDeclaredConstructor(paramTypes));
    }

    /**
     * 处理普通方法
     *
     * @param ctClass    ct类
     * @param paramTypes 参数列表
     * @param function   函数式接口
     */
    @SneakyThrows
    private void dealWithMethod(CtClass ctClass, CtClass[] paramTypes, ProcessFunction<CtClass, CtBehavior> function) {
        CtMethod ctMethod;
        if (paramTypes.length == 0) {
            ctMethod = ctClass.getDeclaredMethod(methodName);
        } else {
            ctMethod = ctClass.getDeclaredMethod(methodName, paramTypes);
        }
        function.process(ctClass, ctMethod);
    }

    /**
     * 获取方法参数类型
     *
     * @param parameterTypes 参数类型列表
     * @return 返回方法参数类型
     */
    private CtClass[] getMethodParamTypes(List<Class<?>> parameterTypes) {
        return parameterTypes.stream().map(this::toCtClass).toArray(CtClass[]::new);
    }

    /**
     * 分离方法参数
     *
     * @param classes 参数列表
     */
    private void detachMethodParams(CtClass[] classes) {
        Optional.ofNullable(classes).ifPresent(v -> Arrays.stream(v).forEach(CtClass::detach));
    }

    /**
     * 转为ctClass
     *
     * @param type 类型
     * @return 返回ctClass
     */
    @SneakyThrows
    private CtClass toCtClass(Class<?> type) {
        return BytecodeUtil.getClassPool().get(type.getName());
    }

    /**
     * 处理函数接口
     *
     * @param <CtClass>  ct类
     * @param <CtBehavior> ct行为
     */
    @FunctionalInterface
    protected interface ProcessFunction<CtClass, CtBehavior> {
        /**
         * 处理
         *
         * @param ctClass  ct类
         * @param ctMethod ct方法
         */
        void process(CtClass ctClass, CtBehavior ctMethod);
    }
}
